#!/usr/bin/python

# Imports
import MySQLdb as mdb
import sys
import os
import random

# DB Connection variable
con = None
# List of lines in the input file
lines = []

if (len(sys.argv) < 3):
    print "Usage: mysql_import_data.py <path to data> <number of records (-1 for all)>"
    sys.exit()
    
# uncompress data
print("\nExtracting data\n");
os.system("gunzip -c %s > ./data.txt" % sys.argv[1]);
print("\nData extracted\n");

# Read in the contents of the data set
with open("./data.txt", 'r') as f:
    # Read in the data
    lines = f.readlines()

# Try to open a connection to the database
try:
    # Connect to the DB
    con = mdb.connect('localhost', 'root', 
        'password', 'mysql_db')
        
	# Grab the cursor
    cursor = con.cursor()

	# drop edges table, drop if it already exists
    sql = "DROP TABLE IF EXISTS mysql_db.facebook_edge"
    cursor.execute(sql);
    sql = "CREATE TABLE IF NOT EXISTS mysql_db.facebook_edge (\
		uid1 INT NOT NULL REFERENCES mysql_db.facebook_node (uid), INDEX(uid1),\
		uid2 INT NOT NULL REFERENCES mysql_db.facebook_node (uid),\
		time TIMESTAMP(8),\
		PRIMARY KEY (uid1, uid2)\
		);"
    cursor.execute(sql)

    # create users table, drop if already exists
    sql = "DROP TABLE IF EXISTS mysql_db.facebook_node"
    cursor.execute(sql);
    sql = "CREATE TABLE IF NOT EXISTS mysql_db.facebook_node (\
	            uid INT NOT NULL PRIMARY KEY, \
                age INT NOT NULL, INDEX(age))"
    cursor.execute(sql)
    

    i = 0
    max_records = int(sys.argv[2])
    print("\nImporting data...\n")
    # Loop over the lines read in
    for line in lines:
        # Split the line
        (uid1, uid2, timestamp) = line.split()
		
		# Check if timestamp is '\N'
        if (timestamp == '\N'):
            # Set timestamp to an empty string
            timestamp = ""

        # Insert users into 
        age = random.randrange(0,100)
        sql = "INSERT IGNORE INTO mysql_db.facebook_node VALUES ('%s','%s')" % (uid1, age)
        cursor.execute(sql);
        age = random.randrange(0,100)
        sql = "INSERT IGNORE INTO mysql_db.facebook_node VALUES ('%s','%s')" % (uid2, age)
        cursor.execute(sql);
	
		# Create the sql statement to insert relationship
        sql = "INSERT INTO mysql_db.facebook_edge VALUES ('%s','%s','%s')" % (uid1, uid2, timestamp)

		# Execute the insert statement
        cursor.execute(sql);

        i = i + 1
        if max_records != -1 and i > max_records:
            break
    print("\nData imported\n");
    
except mdb.Error, e:
    print "Error %d: %s" % (e.args[0], e.args[1])
    sys.exit(1)

finally:
    # Check if we have a connection and close it
    if con:
        con.close()
